#!/usr/bin/env python
"""
:mod:`pynchon.bio.io.fasta`
===========================

Provides read/write support for the fasta file format.
"""
from tempfile import NamedTemporaryFile
from textwrap import wrap
import numpy as np
from pynchon.core.symbol import SymbolVector, SymbolVector0dS1, SymbolArray
from pynchon.bio.aa import AASeq, AAAln
from pynchon.bio.nt import NtSeq


def _get_id_comment(array, key, fval=''):
    """
    """
    id_val = []
    seqs = array.entities if hasattr(array, 'entities') else array
    for seq in seqs:
        seq_id = seq.metadata['header'].split('/')[0].split('|')[0]
        seq_vl = fval
        for line in seq.metadata['comments']:
            if line.startswith(key):
                seq_vl = line.split(':')[1]
                continue
        id_val.append((seq_id, seq_vl))
    return id_val
        
#def fasta_header_parser(ident):
# http://github.com/davisp/bioneb
#    bits = ident.split("|")
#    if len(bits) == 1:
#        return ident
#    ret = {}
#    while len(bits) > 0:
#        itype = bits[0]
#        parts = []
#        for b in bits[1:]:
#            if b in IDENT_TYPES: break
#            parts.append(b)
#        bits = bits[len(parts)+1:]
#        parts = filter(None, parts)
#        if len(parts) == 1:
#            parts = parts[0]
#        if isinstance(ret.get(itype, None), list):
#            ret[itype].append(parts)
#        elif itype in ret:
#            ret[itype] = [ret[itype], parts]
#        else:
#            ret[itype] = parts
#    return ret


def fasta_parser(fh, vector=None, **kwargs):
    """
    Parses a fasta file and generates tuples of (header, sequence, comments).
    If vectors, which should be a ``SymbolVector`` sub-class, is given the 
    generator instead yields it's instances. This parser does not preserve 
    newlines in a fasta file.
    
    Arguments:
    
        - fh (``file``) An open for reading file handle with fasta content or a
          sequence of lines.
        - vector (``SymbolVector`` sub-class) [default =``None``]
    
    Additional keyworded arguments are passed to the constructor of the 
    ``Symbol`` sub-class.
    """
    header = None
    seqs = []
    comments = []
    metabase = kwargs.pop('metadata', {})
    metadata = dict(metabase)
    for line in fh:
        line = line.rstrip()
        if not line:
            continue
        if line[0] == ">":
            if header is not None:
                if vector is None:
                    yield header, "".join(seqs), comments
                else:
                    metadata['header'] = header
                    metadata['comments'] = comments
                    yield vector("".join(seqs), metadata=metadata, **kwargs)        
            header = line[1:].strip()
            seqs = []
            comments = []
            metadata = dict(metabase)
        elif line[0] == ';':
            # unusual comment
            comments.append(line[1:].strip())
        else:
            seqs.append(line)
    if vector is None:
        yield header, "".join(seqs), comments
    else:
        metadata['header'] = header
        metadata['comments'] = comments
        yield vector("".join(seqs), metadata=metadata, **kwargs)

        
def fasta_writer(fh, vector, width=80, **kwargs):
    """
    Writes a ``SymbolVector`` instance or a (header, sequence, comments) 
    sequence to an open for writing file handle in the fasta file format. 
    Supports other  objects with a ``write`` method. The width of the sequence 
    (not comments) can be adjusted with the "width" argument. Symbol vectors
    have to define a 'header' in metadata.
    
    Arguments:
    
        - fh (``file``) An open for writing file handle or object with a 
          ``write`` method.
        - vector (``SymbolVector`` instance or ``sequence``) The contents of 
          the ``sequence`` should be (header, sequence, comments), where header
          and sequence are string and comments is a sequence of strings.
    
    Additional keyworded arguments are passed to the ``textwrap.wrap`` function.
    """
    # get data
    if isinstance(vector, SymbolVector):
        header = vector.metadata['header']
        comments = vector.metadata.get('comments', ())
        seq = "".join(vector.view('S1'))
    else:
        header, seq, comments = vector
    # write data
    lines = [">%s" % header]
    lines.extend([';%s' % line for line in comments])
    lines.extend(wrap(seq, width, break_on_hyphens=False, **kwargs))
    fh.write("\n".join(lines))
    fh.write("\n")
    

def fasta_alignment(fh, vector =None):
    """
    Creates an ``SymbolArray`` object assuming that the contents of the input 
    file are aligned.
    
    Arguments:
    
        - fh (``file``) An open for reading file handle or sequence of lines in
          'fasta' format.
        - vector (``SymbolArray`` sub-class) [default: ``None``]
    
    Additional keyworded arguments are passed to the ``textwrap.wrap`` function.
    """
    parser = fasta_parser(fh, None)
    pieces = list(parser) 
    heds = [p[0] for p in pieces]
    seqs = [p[1] for p in pieces]
    coms = [p[2] for p in pieces]
    # assumed shape of alignment
    shape = (len(seqs), len(seqs[0]))
    assert len(set([len(seq) for seq in seqs])) == 1
    raw_array = np.array("".join(seqs), dtype='c') # create char_array.
    raw_array.shape = shape
    if vector:
        return vector(raw_array, copy=False, fast=True)
    else: 
        return raw_array

def tmp_fasta(vector, **kwargs):
    """
    Writes a ``SymbolVector`` (or ``SymbolArray``) to a temporary file. Returns 
    a file handle open for reading and writing. If the handle is closed or 
    garbage collected the temporary file will be deleted from the file system.
    
    Arguments
    
        - vector (``sequence`` or ``Symbol`` instance). This should be a valid 
          input for ``fasta_writer`` or a ``SymbolArray`` if multiple sequences
          are to be written to one file.
          
    Additional keyworded arguments are passed to ``fasta_writer``.
    """
    tmp_file = NamedTemporaryFile()
    if isinstance(vector, SymbolArray):
        for seq in vector:
            fasta_writer(tmp_file, seq, **kwargs)
    else:
        fasta_writer(tmp_file, vector, **kwargs)
    tmp_file.flush()
    return tmp_file
